import psycopg2
import json
from tqdm import tqdm
import os
from openai import OpenAI

class VBAFile2Database:
    def __init__(self, api_key=None):
        """
        初始化数据库连接
        :param api_key: 阿里云向量嵌入服务的API密钥
        """
        self.DB_CONFIG = {
            "host": "hadoop102",
            "port": 15432,
            "database": "hirain_vba_usage",
            "user": "postgres",
            "password": "postgres"
        }
        # 创建阿里云向量嵌入客户端
        self.embedding_client = OpenAI(
            api_key=api_key,
            base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
        )

    def get_embedding_from_api(self, text):
        """从API获取文本嵌入向量"""
        try:
            completion = self.embedding_client.embeddings.create(
                model="text-embedding-v3",
                input=text,
                dimensions=768,
                encoding_format="float"
            )
            # 从响应中提取向量
            embedding = completion.data[0].embedding
            return embedding
        except Exception as e:
            print(f"获取嵌入向量时出错: {str(e)}")
            return None

    def store_in_db(self, documents, source):
        """存储到PostgreSQL数据库"""
        conn = psycopg2.connect(**self.DB_CONFIG)

        try:
            with conn.cursor() as cur:
                # 创建存储表
                cur.execute("""
                    CREATE TABLE IF NOT EXISTS vba_docs (
                        id SERIAL PRIMARY KEY,
                        content TEXT,
                        vector VECTOR(768),
                        metadata JSONB,
                        collection_name VARCHAR(255)
                    )
                """)

                with tqdm(documents) as pbar:
                    for doc in documents:
                        # 获取文本的向量表示
                        vector = self.get_embedding_from_api(doc['source'] + "::" + doc['answer'])
                        
                        if vector is None:
                            print(f"跳过文档，因为无法获取向量: {doc['source']}")
                            continue

                        metadata = {
                            "source": doc['source'],
                            "collection": doc['collection']
                        }
                        metadata = json.dumps(metadata)
                        
                        cur.execute("""
                        INSERT INTO vba_docs (content, vector, metadata, collection_name)
                        VALUES (%s, %s, %s, %s)
                        """, (doc['answer'], vector, metadata, doc['source']))
                        pbar.update()
                conn.commit()
                print(f"成功插入 {len(documents)} 条记录")
        finally:
            conn.close()
